package com.ruoyi.web.controller;



import com.ruoyi.AI.ChatCompletionMessage;
import com.ruoyi.AI.ChatCompletionRequest;
import com.ruoyi.AI.ChatMessageRole;
import com.ruoyi.AI.Client;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * @author Ryan
 * @version 1.0
 * @date 2024/5/9 15:44
 */
@RequestMapping("/chat")
@RestController
public class ChatSSEController {

    private final ExecutorService executorService = Executors.newFixedThreadPool(10); // 创建固定大小的线程池，可根据并发需求调整

    private final ConcurrentHashMap<String, SseEmitter> emitterMap = new ConcurrentHashMap<>();

    @Autowired
    private Client client;

    @GetMapping(value = "/stream/{userId}", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter streamChat(@RequestParam("message") String userMessage, @PathVariable String userId) {
        if (emitterMap.get(userId) != null) {
            return null;
        }
        SseEmitter emitter = new SseEmitter();
        emitterMap.put(userId, emitter);
        executorService.execute(() -> {
            try {
                getAiReply(userMessage,emitter);
            } catch (Exception e) {
                emitter.completeWithError(e);
                emitterMap.remove(userId);
            } finally {
                // 完成发送后，可以结束SseEmitter
                emitter.complete();
                emitterMap.remove(userId);
            }
        });
        return emitter;
    }

    @GetMapping(value = "/cancel/{userId}")
    public void cancel(@PathVariable String userId) {
        SseEmitter sseEmitter = emitterMap.get(userId);
        if (sseEmitter != null) {
            sseEmitter.complete();
            emitterMap.remove(userId);
        }
    }

    private void getAiReply(String userMessage,SseEmitter emitter) {
        List<ChatCompletionMessage> messages = Arrays.asList(
                new ChatCompletionMessage(ChatMessageRole.USER.value(),
                        userMessage)
        );
        System.out.println("user:"+messages.get(0).getContent());
        try {
            client.ChatCompletionStream(new ChatCompletionRequest(
                    "moonshot-v1-8k",
                    messages,
                    50,
                    0.3f,
                    1
            ),emitter);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }


}
